package com.zhny.test;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.parquet.Strings;
import org.apache.spark.api.java.function.Function;
import org.apache.spark.mllib.linalg.Vector;
import org.apache.spark.mllib.linalg.Vectors;
import scala.Tuple2;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.mllib.regression.LabeledPoint;
import org.apache.spark.mllib.tree.RandomForest;
import org.apache.spark.mllib.tree.model.RandomForestModel;
import org.apache.spark.mllib.util.MLUtils;

//随机森林
public class RandomForestClassificationUtil {
    public static void exc(String dataFilePath, String resultFilePath) {
        SparkConf sparkConf = new SparkConf().setAppName("RandomForestClassification Algorithm").setMaster("local[*]");
        sparkConf.set("spark.driver.allowMultipleContexts", "true");
        JavaSparkContext jsc = new JavaSparkContext(sparkConf);

        // Load and parse the data file.
        String datapath = "";
        if (Strings.isNullOrEmpty(dataFilePath)) {
            datapath = "src/main/resources/data/RFData.txt";
        } else {
            datapath = dataFilePath;
        }

        JavaRDD<LabeledPoint> data = MLUtils.loadLibSVMFile(jsc.sc(), datapath).toJavaRDD();
        // Split the data into training and test sets (30% held out for testing)
        JavaRDD<LabeledPoint>[] splits = data.randomSplit(new double[]{0.7, 0.3});
        JavaRDD<LabeledPoint> trainingData = splits[0];
        JavaRDD<LabeledPoint> testData = splits[1];

        // Train a RandomForest model.
        // Empty categoricalFeaturesInfo indicates all features are continuous.
        int numClasses = 2;
        Map<Integer, Integer> categoricalFeaturesInfo = new HashMap<>();
        int numTrees = 3; // Use more in practice.
        String featureSubsetStrategy = "auto"; // Let the algorithm choose.
        String impurity = "gini";
        int maxDepth = 5;
        int maxBins = 32;
        int seed = 12345;

        RandomForestModel model = RandomForest.trainClassifier(trainingData, numClasses,
                categoricalFeaturesInfo, numTrees, featureSubsetStrategy, impurity, maxDepth, maxBins,
                seed);

        // Evaluate model on test instances and compute test error
        JavaPairRDD<Double, Double> predictionAndLabel =
                testData.mapToPair(p -> new Tuple2<>(model.predict(p.features()), p.label()));
        double testErr =
                predictionAndLabel.filter(pl -> !pl._1().equals(pl._2())).count() / (double) testData.count();
        System.out.println("Test Error: " + testErr);
        System.out.println("Learned classification forest model:\n" + model.toDebugString());

        try {
            FileWriter fos = new FileWriter(new File(resultFilePath));
            fos.write(model.toDebugString());

            Vector v = Vectors.dense(new double[] {3.0, 4.3, 5.5, 6.5, 8.1, 11.5});
            Vector v2 = Vectors.dense(new double[] {3.5, 6.3, 2.5, 9.5, 10.1, 15.5});

            List<Vector> lst = new ArrayList<Vector>();
            lst.add(v);
            lst.add(v2);

            JavaRDD<Vector> features = jsc.parallelize(lst);

            JavaRDD<String> result = features.map(new Function<Vector, String>() {
                @Override
                public String call(Vector arg0) throws Exception {
                    return model.predict(arg0) + arg0.toString();
                }
            });

            fos.write(result.toDebugString());

            fos.flush();
            fos.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        // Save and load model
//        model.save(jsc.sc(), "src/main/resources/data/modal/RandomForestClassificationModel");
//        RandomForestModel sameModel = RandomForestModel.load(jsc.sc(),
//                "src/main/resources/data/modal/RandomForestClassificationModel");

        jsc.stop();
    }
}
